// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/common/caching/cache_ttl_controller.h>

#include <pollux/common/caching/async_data_cache.h>

namespace kumo::pollux::cache {
    std::unique_ptr<CacheTTLController> CacheTTLController::instance_ = nullptr;

    bool CacheTTLController::addOpenFileInfo(
        uint64_t fileNum,
        int64_t openTimeSec) {
        auto lockedFileMap = fileInfoMap_.wlock();
        auto it = lockedFileMap->find(fileNum);
        if (it == lockedFileMap->end() || it->second.removeInProgress) {
            lockedFileMap->insert_or_assign(fileNum, RawFileInfo{openTimeSec, false});
            return true;
        }
        return false;
    }

    CacheAgeStats CacheTTLController::getCacheAgeStats() const {
        auto lockedFileMap = fileInfoMap_.rlock();

        if (lockedFileMap->empty()) {
            return CacheAgeStats{.maxAgeSecs = 0};
        }

        // Use the oldest file open time to calculate the max possible age of cache
        // entries loaded from the files.
        int64_t minOpenTime = std::numeric_limits<int64_t>::max();
        for (auto it = lockedFileMap->cbegin(); it != lockedFileMap->cend(); it++) {
            minOpenTime = std::min<int64_t>(minOpenTime, it->second.openTimeSec);
        }

        int64_t maxAge = getCurrentTimeSec() - minOpenTime;
        return CacheAgeStats{.maxAgeSecs = std::max<int64_t>(maxAge, 0)};
    }

    void CacheTTLController::applyTTL(int64_t ttlSecs) {
        int64_t maxOpenTime = getCurrentTimeSec() - ttlSecs;

        melon::F14FastSet<uint64_t> filesToRemove =
                getAndMarkAgedOutFiles(maxOpenTime);
        if (filesToRemove.empty()) {
            KLOG(INFO) << "No cache entry is out of TTL " << ttlSecs << ".";
            return;
        }

        melon::F14FastSet<uint64_t> filesRetained;
        bool success = cache_.removeFileEntries(filesToRemove, filesRetained);

        KLOG(INFO) << (success ? "Succeeded" : "Failed") << " applying cache TTL of "
                << ttlSecs << " seconds. Entries from " << filesToRemove.size()
                << " files are to be removed, while " << filesRetained.size()
                << " files are retained";
        if (success) {
            cleanUp(filesRetained);
        } else {
            reset();
        }
    }

    melon::F14FastSet<uint64_t> CacheTTLController::getAndMarkAgedOutFiles(
        int64_t maxOpenTimeSecs) {
        auto lockedFileMap = fileInfoMap_.wlock();

        melon::F14FastSet<uint64_t> fileNums;

        for (auto it = lockedFileMap->begin(); it != lockedFileMap->end(); it++) {
            if (it->second.removeInProgress ||
                it->second.openTimeSec < maxOpenTimeSecs) {
                fileNums.insert(it->first);
                it->second.removeInProgress = true;
            }
        }

        return fileNums;
    }

    void CacheTTLController::cleanUp(
        const melon::F14FastSet<uint64_t> &filesToRetain) {
        fileInfoMap_.withWLock([&](auto &fileMap) {
            auto it = fileMap.begin();
            while (it != fileMap.end()) {
                if (!it->second.removeInProgress) {
                    it++;
                    continue;
                }
                if (filesToRetain.count(it->first) > 0) {
                    it->second.removeInProgress = false;
                    it++;
                    continue;
                }
                it = fileMap.erase(it);
            }
        });
    }

    void CacheTTLController::reset() {
        fileInfoMap_.withWLock([](auto &fileMap) {
            for (auto &[_, fileInfo]: fileMap) {
                fileInfo.removeInProgress = false;
            }
        });
    }
} // namespace kumo::pollux::cache
